
"""
Dynamic Subgroup Pruning (Paper Section 2.3)
Implements Equations 5-8: Spectral clustering with membership pruning
"""

import numpy as np
from sklearn.cluster import SpectralClustering

class SpectralSubgroupOptimizer:
    def __init__(self, k=5, cutoff_dim=20, prune_thresh=1.5):
        """
        Args:
            k: Initial subcluster count
            cutoff_dim: Truncated SVD dimension
            prune_thresh: Noise pruning threshold (σ in Equation 7)
        """
        self.k = k
        self.cutoff_dim = cutoff_dim
        self.prune_thresh = prune_thresh

    def decompose_kernel_matrix(self, adj_matrix):
        """Kernel Matrix Spectral Decomposition (Equation 5-6)"""
        # Truncated SVD
        U, S, Vt = np.linalg.svd(adj_matrix)
        var_retained = np.cumsum(S**2 / (S**2).sum())
        k = np.argmax(var_retained >= 0.9) + 1  # 90% variance criteria
        self.cutoff_dim = min(k, self.cutoff_dim)
        
        return U[:, :self.cutoff_dim] @ np.diag(S[:self.cutoff_dim])

    def perform_clustering(self, decomposed_matrix):
        """Spectral Clustering with Dynamic Pruning"""
        # Initial clustering
        clustering = SpectralClustering(n_clusters=self.k, affinity='nearest_neighbors')
        raw_labels = clustering.fit_predict(decomposed_matrix)
        
        # Membership pruning (Equation 7-8)
        pruned_labels = []
        cluster_centers = [np.mean(decomposed_matrix[raw_labels==i], axis=0) 
                          for i in range(self.k)]
        
        for i, vec in enumerate(decomposed_matrix):
            dists = [np.linalg.norm(vec - center) for center in cluster_centers]
            min_dist = np.min(dists)
            if min_dist < self.prune_thresh:
                pruned_labels.append(np.argmin(dists))
            else:
                pruned_labels.append(-1)  # Mark as noise
                
        return np.array(pruned_labels)
